#include <winrt/Windows.Foundation.h>
#include <winrt/Windows.Graphics.Capture.h>
#include <winrt/Windows.Graphics.DirectX.Direct3D11.h>
#include <d3d11.h>
#include <dxgi1_2.h>

// WinRT interop 头文件（关键）
#include <windows.graphics.directx.direct3d11.interop.h>
#include <windows.graphics.capture.interop.h>
#include <winrt/base.h>

#include <iostream>

inline void ThrowIfFailed(HRESULT hr)
{
    if (FAILED(hr))
    {
        throw winrt::hresult_error(hr);
    }
}

// 辅助函数：从 HWND 创建 CaptureItem
winrt::Windows::Graphics::Capture::GraphicsCaptureItem CreateCaptureItemFromWindow(HWND hwnd)
{
    winrt::Windows::Graphics::Capture::GraphicsCaptureItem item{ nullptr };

    auto factory = winrt::get_activation_factory<
        winrt::Windows::Graphics::Capture::GraphicsCaptureItem,
        IGraphicsCaptureItemInterop>();

    winrt::guid iid = winrt::guid_of<winrt::Windows::Graphics::Capture::GraphicsCaptureItem>();
    ThrowIfFailed(factory->CreateForWindow(hwnd, iid, reinterpret_cast<void**>(winrt::put_abi(item))));

    return item;
}

// 辅助函数：创建 IDirect3DDevice
winrt::Windows::Graphics::DirectX::Direct3D11::IDirect3DDevice CreateD3DDevice()
{
    winrt::com_ptr<ID3D11Device> d3dDevice;
    D3D_FEATURE_LEVEL levels[] = { D3D_FEATURE_LEVEL_11_0 };
    ThrowIfFailed(D3D11CreateDevice(
        nullptr,
        D3D_DRIVER_TYPE_HARDWARE,
        nullptr,
        D3D11_CREATE_DEVICE_BGRA_SUPPORT,
        levels,
        ARRAYSIZE(levels),
        D3D11_SDK_VERSION,
        d3dDevice.put(),
        nullptr,
        nullptr));

    auto dxgiDevice = d3dDevice.as<IDXGIDevice>();

    winrt::com_ptr<::IInspectable> inspectable;
    ThrowIfFailed(CreateDirect3D11DeviceFromDXGIDevice(dxgiDevice.get(), inspectable.put()));

    return inspectable.as<winrt::Windows::Graphics::DirectX::Direct3D11::IDirect3DDevice>();
}

static int frames = 0;

int main()
{
    winrt::init_apartment();

    // 找到记事本窗口
    HWND hwnd = FindWindowW(L"Notepad", nullptr);
    if (!hwnd)
    {
        std::cout << "Please open Notepad first!\n";
        return 1;
    }

    auto item = CreateCaptureItemFromWindow(hwnd);
    auto d3dDevice = CreateD3DDevice();

    auto size = item.Size();

    auto framePool = winrt::Windows::Graphics::Capture::Direct3D11CaptureFramePool::Create(
        d3dDevice,
        winrt::Windows::Graphics::DirectX::DirectXPixelFormat::B8G8R8A8UIntNormalized,
        2,
        size);

    auto session = framePool.CreateCaptureSession(item);

    framePool.FrameArrived([&](auto& sender, auto&)
        {
            auto frame = sender.TryGetNextFrame();
            if (frame)
            {
                auto contentSize = frame.ContentSize();
                std::cout << "Captured a frame: "
                    << contentSize.Width << "x" << contentSize.Height << ", frames: " << frames << "\n";
                frames++;
            }
        });

    session.StartCapture();

    MSG msg;
    while (GetMessage(&msg, nullptr, 0, 0))
    {
        TranslateMessage(&msg);
        DispatchMessage(&msg);
    }


    std::cout << "Capturing window, press Enter to exit...\n";
    std::wstring dummy;
    std::getline(std::wcin, dummy);

    return 0;
}
